# coding : UTF-8
"""
作者：BingBO   时间：2022年08月11日
自动调整代码格式 ：Alt+Ctrl+L
"""
import cv2 as cv
import numpy as np
import glob
import math
from matplotlib import pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import calibrate_camera as cc
import extract_laser_from_checkerboard


# 激光平面标定步骤：
# 1.返回棋盘格中的单线激光 直线拟合参数s1
# 2.求出棋盘格的9列内角点的 直线拟合参数 l1,l2,l3,...,l9
# 3.计算s1和l1,l2,l3,...,l9的9个图像交点
# 4.根据标定参数，将图像交点转换到相机坐标系下的3D坐标
# 5.拟合所有这些激光点的平面参数(参考相机坐标系)


# 获取标定棋盘格每一列与单线激光的多个交点（图像2D坐标）
def get_intersection_coordinates(laser_in_checkerboard_list, imagePoints):
    # 按照和标定时的遍历顺序一致去
    laser_k_b_list = []
    for frame_path in laser_in_checkerboard_list:
        laser_frame = cv.imread(frame_path)
        laser_frame = cv.resize(laser_frame, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)

        # 返回棋盘格中的单线激光 直线拟合参数s1
        k, b = extract_laser_from_checkerboard.get_laser_polyfit_from_checkerboard(laser_frame)
        laser_k_b_list.append([k, b])

    # 定位在图像中的位置（列，行）
    each_checkerboard_list = []
    for imagePoint in imagePoints:
        each_col_k_b_list = []
        for i in range(9):
            # 存储每一列的坐标值
            col_coordis = imagePoint[i::9]  # 步长为9
            # 对每一列的坐标值 进行直线拟合
            x_list = []
            y_list = []
            for item in col_coordis:
                x_list.append(item[0][0])
                y_list.append(item[0][1])
            print("x_list :", x_list)
            print("y_list :", y_list)
            k, b = np.polyfit(x_list, y_list, 1)  # 用最小二乘法（含COG_V误差）拟合得到的直线参数
            each_col_k_b_list.append([k, b])
        each_checkerboard_list.append(each_col_k_b_list)

    # 计算每张图的X个交点图像坐标（不进行CCS转换）
    all_points_per_checkerboard = []
    for idx in range(len(laser_k_b_list)):
        nine_points_of_each = []  # 保存当前棋盘的9个交点

        laser_k_b = laser_k_b_list[idx]
        col_lines_k_b = each_checkerboard_list[idx]  # 9列线的kb参数
        for each_k_b in col_lines_k_b:
            laser_k = laser_k_b[0]
            laser_b = laser_k_b[1]
            col_k = each_k_b[0]
            col_b = each_k_b[1]
            # 两条一般直线（排除任何一条直线垂直的情况）
            x = (col_b - laser_b) / (laser_k - col_k)
            y = laser_k * x + laser_b
            # y坐标的范围限制在图形内部
            if y > 0:
                nine_points_of_each.append([round(x), round(y)])
        all_points_per_checkerboard.append(nine_points_of_each)
    return all_points_per_checkerboard  # 图像坐标


# 将每张图上的9个图像点，通过每张图的内外参数，转换为相机坐标3D点
def image_to_camera_coordinates(all_points_per_checkerboard):
    # 读取标定参数
    cameraMatrix, distCoeffs, R, T, imagePoints = cc.load_calibrate_data()
    # 先判断一下参数对数是否匹配
    # 保存所有的激光平面的点
    all_ccs_points = []
    for idx in range(len(all_points_per_checkerboard)):
        R_Matrix = cv.Rodrigues(R[idx].ravel())  # 每张棋盘图的旋转向量
        T_Vecs = T[idx]  # 每张棋盘图的平移向量

        # 外参矩阵 M1 (世界坐标系WCS to 相机坐标系CCS)
        # 外参 ： 旋转矩阵R_Matrix[0] 平移向量tVecs[i]
        R_T = np.hstack((R_Matrix[0], T_Vecs))
        M1 = np.vstack((R_T, np.array([0, 0, 0, 1])))
        print("M1:\n", M1)

        #  (相机坐标系下) 棋盘靶标平面系数方程系数
        inv_M1 = np.linalg.inv(M1)
        a1, b1, c1, d1 = inv_M1[2][0], inv_M1[2][1], inv_M1[2][2], inv_M1[2][3]
        # print("(相机坐标系下)棋盘靶标平面a1*x+b1*y+c1*z+d1=0系数: [a1, b1, c1, d1]:", inv_M1[2])

        #  (棋盘靶标平面 + 激光条纹上)的 图像坐标的相机坐标点
        fx, fy, Cx, Cy = cameraMatrix[0][0], cameraMatrix[1][1], cameraMatrix[0][2], cameraMatrix[1][2]

        # 取出每张棋盘图的9个交点的坐标
        nine_points_per_checkerboard = all_points_per_checkerboard[idx]
        for each_point in nine_points_per_checkerboard:
            # 求去畸变图像坐标
            undistort_point = cv.undistortPoints(tuple(each_point), cameraMatrix, distCoeffs, P=cameraMatrix)
            u = undistort_point[0][0][0]
            v = undistort_point[0][0][1]

            # 计算相机坐标
            Zc = -d1 / (a1 * (u - Cx) / fx + b1 * (v - Cy) / fy + c1)
            Yc = Zc * (v - Cy) / fy
            Xc = Zc * (u - Cx) / fx
            all_ccs_points.append([Xc, Yc, Zc])
            # print("棋盘激光点(u%d,v%d) 反向映射的Camera坐标为：" % ((j + 1), (j + 1)), (Xc, Yc, Zc))
    return all_ccs_points


def save_laser_ccs_points_data(laser1_ccs_points, laser2_ccs_points):
    # 保存所有图像点文件
    f = open('./result/激光平面标定CCS点.txt', 'w+')
    f.write("laser1_ccs_points:\n" + str(np.array(laser1_ccs_points)) + '\n')
    f.write("****************************************************\n")
    f.write("laser2_ccs_points:\n" + str(np.array(laser2_ccs_points)) + '\n')
    f.close()

    np.savez('./result/laser_ccs_points_data.npz', laser1_ccs_points=laser1_ccs_points,
             laser2_ccs_points=laser2_ccs_points)
    print('laser_ccs_points_data.npz文件已保存')


def load_laser_ccs_points_data():
    with np.load('./result/laser_ccs_points_data.npz') as X:
        laser1_ccs_points, laser2_ccs_points = [X[i] for i in ('laser1_ccs_points', 'laser2_ccs_points')]
    return laser1_ccs_points, laser2_ccs_points


# 激光平面拟合
def laser_plane_fit(all_ccs_points):
    # all_ccs_points = np.array(all_ccs_points)  # 列表转换成array
    all_ccs_points = all_ccs_points.reshape(-1, 3)  # 转换成n行3列
    row = all_ccs_points.shape[0]
    x, y, z = all_ccs_points[:, 0], all_ccs_points[:, 1], all_ccs_points[:, 2]

    # 创建系数矩阵A、向量b
    A = np.ones((row, 3))
    b = np.zeros((row, 1))
    for i in range(0, row):
        A[i, 0] = x[i]
        A[i, 1] = y[i]
        b[i, 0] = z[i]

    # 通过 X=(AT*A)-1*AT*b 求解超定方程
    A1 = np.linalg.inv(np.dot(A.T, A))
    A2 = np.dot(A1, A.T)
    X = np.dot(A2, b)
    # 激光平面方程 z = ax + by + c
    a, b, c = X[0, 0], X[1, 0], X[2, 0]
    print('激光平面拟合结果为：Zc = %.3f * Xc + %.3f * Yc + %.3f' % (a, b, c))

    # 计算所有点到拟合平面的总距离差d
    d = 0
    for i in range(0, row):
        d += math.sqrt((a * x[i] + b * y[i] + c - z[i]) * (a * x[i] + b * y[i] + c - z[i]))
    error = d / row
    print('平均每个点到拟合激光平面的距离误差：%.*f' % (3, error))
    return a, b, c


# 显示相机坐标系下的拟合平面
def laser_plane_display():
    laser1_ccs_points, laser2_ccs_points = load_laser_ccs_points_data()

    a1, b1, c1 = laser_plane_fit(laser1_ccs_points)
    a2, b2, c2 = laser_plane_fit(laser2_ccs_points)
    # list转成np array，才能执行切片操作
    laser1_ccs_points = np.array(laser1_ccs_points)
    laser2_ccs_points = np.array(laser2_ccs_points)
    x1, y1, z1 = laser1_ccs_points[:, 0], laser1_ccs_points[:, 1], laser1_ccs_points[:, 2]
    x2, y2, z2 = laser2_ccs_points[:, 0], laser2_ccs_points[:, 1], laser2_ccs_points[:, 2]

    # 展示图像
    fig1 = plt.figure()  # 定义新的三维坐标轴
    # ax1 = Axes3D(fig)  # 创建3d图形的两种方式
    axis1 = fig1.add_subplot(1, 1, 1, projection='3d')  # 将这块画布分为1×1，然后1对应的就是1号区  并指定为3维

    plt.title("laser_plane")
    axis1.set_xlabel('X', fontsize=14)
    axis1.set_ylabel('Y', fontsize=14)
    axis1.set_zlabel('Z', fontsize=14)
    axis1.scatter(x1, y1, z1, c='blue', marker='o')  # 生成散点图
    axis1.scatter(x2, y2, z2, c='blue', marker='o')

    # 平面各个坐标的范围
    xx = np.arange(-80, 100, 10)
    yy = np.arange(-80, 100, 10)
    x_plane, y_plane = np.meshgrid(xx, yy)
    z1_plane = a1 * x_plane + b1 * y_plane + c1
    z2_plane = a2 * x_plane + b2 * y_plane + c2

    # 画出相机坐标系
    x, y, z = [0, 50], [0, 0], [0, 0]
    axis1.text(50, 0, 0, "Cx")
    axis1.plot(x, y, z, color='blue', label="X")

    x, y, z = [0, 0], [0, 50], [0, 0]
    axis1.text(0, 50, 0, "Cy")
    axis1.plot(x, y, z, color='green', label="Y")

    x, y, z = [0, 0], [0, 0], [0, 50]
    axis1.text(0, 0, 50, "Cz")
    axis1.plot(x, y, z, color='red', label="Z")

    # 绘制平面或曲面图
    # rstride: 行采样之间的跨度     cstride: 列采样之间的跨度
    # alpha 透明度
    axis1.plot_surface(x_plane, y_plane, z1_plane, label="laser1", alpha=0.5, rstride=1, cstride=1, color='Red')
    axis1.plot_surface(x_plane, y_plane, z2_plane, label="laser2", alpha=0.5, rstride=1, cstride=1, color='Red')
    axis1.set_zlim(0, 200)

    plt.show()


if __name__ == "__main__":
    t1 = cv.getTickCount()

    # checkerboard_image_path = r"E:\SHU\Research Group\LaserVisionSensor\Calibrate\checkerboard\*.jpg"
    # checkerboard_images = glob.glob(checkerboard_image_path)
    # checkerboard_images.sort(key=len)  # 排序
    #
    # checkerboard_path_list = checkerboard_images[0::3]  # 棋盘格
    # laser1_path_list = checkerboard_images[1::3]  # 棋盘格 + laser1
    # laser2_path_list = checkerboard_images[2::3]  # 棋盘格 + laser2
    #
    # cameraMatrix1, distCoeffs1, R, T, imagePoints = cc.load_calibrate_data()
    # laser1_all_points_per_checkerboard = get_intersection_coordinates(laser1_path_list, imagePoints)
    # laser2_all_points_per_checkerboard = get_intersection_coordinates(laser2_path_list, imagePoints)
    #
    # laser1_all_ccs_points = image_to_camera_coordinates(laser1_all_points_per_checkerboard)
    # laser2_all_ccs_points = image_to_camera_coordinates(laser2_all_points_per_checkerboard)
    # save_laser_ccs_points_data(laser1_all_ccs_points, laser2_all_ccs_points)

    laser_plane_display()

    t2 = cv.getTickCount()
    print("process:%s ms" % ((t2 - t1) / cv.getTickFrequency() * 1000))

    cv.waitKey(0)
    cv.destroyAllWindows()
